import termcolor
import atexit
import json
from .tts import play_audio_text
from .safety import get_safety_timeout
from .core import ask_question

def pretty_clear_screen():
    print()
    print(termcolor.colored('---------------', 'light_grey'))
    print()

def pretty_print_embedding(nfrom, nto, dt):
    msg = termcolor.colored(f'{nfrom} -> {nto} messages', 'red')
    if dt is not None:
        msg += termcolor.colored(f" ({dt:.2f}s)", 'light_grey')
    print(msg)

def pretty_print_message(message, dt=None, append=True, finish=True):
    role_to_color = {
        "system": "red",
        "user": "green",
        "assistant": "blue",
        "tool": "magenta",
    }
    if not isinstance(message, dict):
        message = message.model_dump()
    msg = None
    if message["role"] == "system":
        msg = f"system: {message['content']}"
    elif message["role"] == "user":
        msg = f"user: {message['content']}"
    elif message["role"] == "assistant":
        if message.get("tool_calls"):
            msg = f"assistant: {message['tool_calls']}"
        if message.get("content"):
            msg = f"assistant: {message['content']}"
    elif message["role"] == "tool":
        msg = f"function ({message['name']}): {message['content']}"
    if msg is None:
        return
    msg = termcolor.colored(msg, role_to_color[message["role"]]) # , None, ['bold'] if bold else None)
    if dt is not None:
        msg += termcolor.colored(f" ({dt:.2f}s)", 'light_grey')
    print(msg)

def function_safety_wait(name, args):
    print(termcolor.colored(name, 'magenta'), termcolor.colored(json.dumps(args, ensure_ascii=False, indent=4), 'light_magenta'))
    timeout = get_safety_timeout(name, args)
    if timeout <= 0:
        return True
    try:
        from inputimeout import inputimeout, TimeoutOccurred
    except ImportError:
        inputimeout = lambda prompt="", timeout=None: input(prompt)
        TimeoutOccurred = StopIteration
    try:
        inputimeout(termcolor.colored('[Enter to pause]', 'light_grey'), timeout)
    except KeyboardInterrupt:
        print()
        return False
    except TimeoutOccurred:
        return True
    except EOFError:
        return False
    try:
        choice = input(termcolor.colored('Continue action [y/N]?', 'light_magenta'))
    except KeyboardInterrupt:
        print()
        return False
    except EOFError:
        return False
    return choice.lower().strip() == 'y'

def question_repl(model):
    print(termcolor.colored('Welcome to Assistant!', 'light_yellow'))
    print(termcolor.colored('Press /? for help', 'light_grey'))
    import readline
    readline.read_history_file()
    atexit.register(readline.write_history_file)
    options = {
        'temperature': 0.0,
        'max_tokens': 500,
    }
    callbacks = {
        'message': pretty_print_message,
        'answer': play_audio_text,
        'check': function_safety_wait,
    }
    messages = []
    while True:
        try:
            prompt = input('>>> ')
        except KeyboardInterrupt:
            print()
            continue
        except EOFError:
            break
        prompt = prompt.strip()
        if not prompt:
            continue
        elif prompt == '/r':
            messages = []
            pretty_clear_screen()
            print('(New context created)')
            continue
        elif prompt == '/3':
            model = 'gpt-3.5-turbo-0125'
            print('(Switched to GPT-3.5)')
            continue
        elif prompt == '/p':
            while len(messages) and messages[-1]['role'] != 'user':
                messages.pop()
            if len(messages) == 0:
                continue
            prompt = messages[-1]['content']
        elif prompt == '/1':
            options['temperature'] = 1.0
            print('(Temperature: 1.0)')
            continue
        elif prompt == '/0':
            options['temperature'] = 0.0
            print('(Temperature: 0.0)')
            continue
        elif prompt == '/4':
            model = 'gpt-4-turbo-preview'
            print('(Switched to GPT-4)')
            continue
        elif prompt == '/q':
            break
        elif prompt == '/?':
            print(termcolor.colored('''
/? - show this help message
/q - quit the program
/0 - set temperature to 0
/1 - set temperature to 1
/3 - switch to gpt-3.5
/4 - switch to gpt-4
/r - reset a new context
/p - regenerate last answer
''', 'light_cyan'))
            continue
        elif prompt.startswith('/'):
            continue
        import readline
        readline.add_history(prompt)
        message = {
            'role': 'user',
            'content': prompt,
        }
        pretty_print_message(message)
        messages.append(message)
        answer = ask_question(model, messages, options, callbacks)
        if answer:
            print(termcolor.colored('>>> ' + answer, 'light_yellow'))

if __name__ == '__main__':
    question_repl('gpt-3.5-turbo-0125')
